#include <stdio.h>

#include "stdinc.h"
#include "ntundoc.h"
#include <winnt.h>


void FormatLastError2()
{
		DWORD lastErr = GetLastError();
		CHAR errBuff[1024];
		FormatMessage(FORMAT_MESSAGE_FROM_SYSTEM,NULL,lastErr,0,errBuff,1024,(va_list*)NULL);
		MessageBox(NULL,errBuff,"Error!",MB_OK);
}


int FindLauncherFunctions(PROCESS_INFORMATION *p, DWORD *GPA,DWORD *LLA)
{
	typedef LONG (WINAPI NTQIP)(HANDLE, PROCESSINFOCLASS, PVOID, ULONG, PULONG);
	NTQIP *lpfnNtQueryInformationProcess;

	PROCESS_BASIC_INFORMATION pbi;
	pbi.PebBaseAddress = (_PEB *) 0x7ffdf000;

	DWORD dwSize;

	HMODULE hLibrary = GetModuleHandle(_T("ntdll.dll"));
	if (NULL != hLibrary) 
	{
		lpfnNtQueryInformationProcess = (NTQIP *) GetProcAddress(hLibrary, 
											 "NtQueryInformationProcess");
		if (NULL != lpfnNtQueryInformationProcess)
			(*lpfnNtQueryInformationProcess)(p->hProcess, 
				 ProcessBasicInformation, &pbi, sizeof(pbi), &dwSize);
	}
	 char buf[2048];
	
	DWORD oldProtect;
	SIZE_T bWritten;
	BOOL result = VirtualProtectEx(p->hProcess,(LPVOID)pbi.PebBaseAddress,sizeof(PEB),PAGE_READWRITE,&oldProtect);
	if (result <0)
	{
		FormatLastError2();
		return -1;
	}
	PEB peb;
	result = ReadProcessMemory(p->hProcess,pbi.PebBaseAddress,&peb,sizeof(PEB),&bWritten);
	if (result <0)
	{
		FormatLastError2();
		return -1;
	}
	_IMAGE_DOS_HEADER idh;
	result = ReadProcessMemory(p->hProcess,peb.ImageBaseAddress,&idh,sizeof(_IMAGE_DOS_HEADER),&bWritten);
	if (result <0)
	{
		FormatLastError2();
		return -1;
	}
	_IMAGE_NT_HEADERS inh;
	_IMAGE_NT_HEADERS *pinh = (_IMAGE_NT_HEADERS*)((char*)peb.ImageBaseAddress+idh.e_lfanew);

	result = ReadProcessMemory(p->hProcess,pinh,&inh,sizeof(_IMAGE_NT_HEADERS),&bWritten);
	if (result <0)
	{
		FormatLastError2();
		return -1;
	}
	DWORD add;
	add = inh.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress;
	_IMAGE_IMPORT_DESCRIPTOR *pidesc = (_IMAGE_IMPORT_DESCRIPTOR*)
									   ((char*)peb.ImageBaseAddress+
										inh.OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT].VirtualAddress);
	bool done = false;
	_IMAGE_IMPORT_DESCRIPTOR importDescriptor;
	while (!done)
	{
		
		result = ReadProcessMemory(p->hProcess,pidesc,&importDescriptor,sizeof(_IMAGE_IMPORT_DESCRIPTOR),&bWritten);
		if (result <0)
		{
			FormatLastError2();
			return -1;
		}
		if (importDescriptor.Characteristics == 0) 
			break;
	
		
		char *pname = (char*)peb.ImageBaseAddress+importDescriptor.Name;
		result = ReadProcessMemory(p->hProcess,pname,buf,256,&bWritten);
		if (result <0)
		{
			FormatLastError2();
			return -1;
		}
		
		if (_stricmp(buf,"kernel32.dll") == 0)
		{
			bool GPAFound=false,LLAFound=false;
			IMAGE_THUNK_DATA *pitd = (IMAGE_THUNK_DATA*)
									 ((char*)peb.ImageBaseAddress+
									  importDescriptor.FirstThunk);
			IMAGE_THUNK_DATA itd;
			while (!GPAFound || !LLAFound)
			{
				result = ReadProcessMemory(p->hProcess,pitd,&itd,sizeof(IMAGE_THUNK_DATA),&bWritten);
				if (result <0)
				{
					FormatLastError2();
					return -1;
				}
				if (itd.u1.AddressOfData == 0)
					break;
				char *fname = (char*)peb.ImageBaseAddress+itd.u1.AddressOfData + 2; // +2 for skipping a 2 byte hint
				result = ReadProcessMemory(p->hProcess,fname,&buf,256,&bWritten);
				if (result <0)
				{
					FormatLastError2();
					return -1;
				}
				if (!GPAFound && _stricmp(buf,"GetProcAddress") == 0)
				{
					GPAFound = true;
					*GPA = (DWORD)pitd;
				}
				if (!LLAFound && _stricmp(buf,"LoadLibraryA") == 0)
				{
					LLAFound = true;
					*LLA = (DWORD)pitd;
				}
				pitd++;
			}
			if (!LLAFound || !GPAFound)
			{
				MessageBox(NULL,"GetProcAddress or LoadLibrary have not been found, aborting","Error",MB_OK);
				return -1;
			}
			done = true;
		}
		pidesc++;
	};
	if (!done)
	{
		MessageBox(NULL,"Kernel32 not found in file, aborting","Error",MB_OK);
		return -1;
	}
	return 1;
	
}

